--
-- Copyright Benjamin Ylvisaker, 2012.  License information to come later.
-- For now it's not licensed for any use.
-- 
-- Take some register code and turn it into some queue code... easy as that
--

-- idea for sleeping value algorithm:
-- start with each reg/BBL pair individually.  If the reg is live but 
-- not accessed, it might be interesting

{-# LANGUAGE OverloadedStrings #-}

import Control.Exception (assert)
-- import Data.List
import Data.Word                         as W
import Data.List                         as L
-- import Data.Bits
-- import Data.Ord
import Data.Graph.Inductive.Dot          as Dot
import qualified System.IO               as SIO
-- import qualified System.Cmd              as Cmd
import qualified Debug.Trace             as T
import qualified Data.Maybe              as MB
import qualified Data.Sequence           as Q
import Data.Sequence ((<|),(|>),(><),ViewL ((:<)))
-- import qualified Data.Set                as S
import qualified Data.Map                as M
-- import qualified Data.Foldable           as F
import qualified Data.IntMap             as IM
-- import qualified Data.IntSet             as IS
import qualified Text.JSON               as JS
-- import qualified Data.Aeson              as AE
-- import qualified Data.HashMap.Strict     as HM
import qualified Data.Graph.Inductive    as G
import Data.Graph.Inductive ((&))
import System.Random                     as R
import qualified System.Environment      as Env
-- import qualified Data.GraphViz as GV

import Miscellaneous_utilities
import Graph_utilities

-- type IntSeq = Q.Seq Int




-- New shoes!  Here's the plan:
-- A. Make a schedule that respects real dependencies, but not queue contraints
--    by selecting the next instruction uniformly at random from the ready pool.
-- B. Estimate the cost of queue instructions thusly:
--   1. Fanout = 0 => 1 unit (can actually be done before scheduling; whatever)
--   2. Fanout > 1 => Fanout - 1 units (can be done before scheduling)
--   3. Figure out the lifetime overlapping/nesting of r-live ranges.  Each
--      nested pair is 1 unit.
--     a. Unassigned live-ins/-outs are a bit of a pain.  Here's an idea:  The
--        instructions get to pick which live-i/o is which in the order they
--        appear in the schedule.  There will be a little bit of hackery to get
--        this right, but not much.
-- C. Do a local optimization like scanning back and forth along the schedule
--    evaluating swapping adjacent pairs of instructions.  If the swap lowers
--    the overall cost estimate, use it, otherwise don't.  This will be at worst
--    N^2, but N is usually pretty small, and I bet the number of iterations
--    will be quite small on average.

-- |random_schedule generates a schedule for the instructions in g.  The
-- |schedule respects all the dependencies in g, but is not necessarily queue-
-- |legal, and will not contain any queue manipulation instructions.  I'm pretty
-- |sure the schedule is selected uniformly at random from all legal schedules
-- |of g.
random_schedule :: R.RandomGen r => Dfg -> r -> [G.Node]
random_schedule dfg_init rgen = random_one_instr dfg_init initial_ready_pool rgen
  where
    is_ready_net g (_,net_id) =
        case g_unsafe_match_net "gimme input net" net_id g of
          ([], _, _, _) -> True
          ([(_,live_in_maybe)], _, _, _) -> live_in_maybe == live_in_hub_id
          _ -> False
    is_ready_instr g n_id =
        case g_unsafe_match "init red" n_id g of
          ((preds, _, DfgOp _, _), _) -> all (is_ready_net g) preds
          ((preds, _, DfgDrop, _), _) -> all (is_ready_net g) preds
          _ -> False
    initial_ready_pool = Q.fromList (filter (is_ready_instr dfg_init) (G.nodes dfg_init))

    random_one_instr :: R.RandomGen r => Dfg -> Q.Seq G.Node -> r -> [G.Node]
    -- random_one_instr _ _ _ | T.trace ("one rand!:") False = undefined
    random_one_instr dfg_randing ready_prev r_gen_prev =
        if Q.null ready_prev then
            []
        else
            {--T.trace "rand one"--} n_instr : (random_one_instr dfg_one ready_next r_gen_next)
            where
              -- Get an instr
              n_instr = Q.index ready_prev random_index
              (random_index, r_gen_next) = R.randomR (0, pool_sz - 1) r_gen_prev
              pool_sz = Q.length ready_prev
              -- Get the next ready pool
              ready_new_raw = foldl follow_succ [] succs
              ready_next = ready_without_n >< (Q.fromList $ nub $ sort ready_new_raw)
              ready_without_n = Q.take random_index       ready_prev
                             >< Q.drop (random_index + 1) ready_prev
              ((_,_,_,succs), dfg_one) =
                  --T.trace ("picked a random "++show n_instr++"  "++show ready_prev)
                  g_unsafe_match ("rrr "++show n_instr) n_instr dfg_randing
              follow_succ ready_following (l,successor) =
                  if is_ready_net dfg_one (l,successor) then
                      foldl follow_net ready_following (G.suc dfg_one successor)
                  else
                      ready_following
              follow_net ready_netting next_instr =
                  if is_ready_instr dfg_one next_instr then
                      next_instr : ready_netting
                  else
                      ready_netting

random_schedule_try_a_few :: R.RandomGen r => Dfg -> r -> [(Int, [G.Node])]
random_schedule_try_a_few dfg r_gen_init = try_one 1 r_gen_init []
  where
    num_to_try = (10 :: Integer)
    try_one counter r_gen schedules =
        if counter > num_to_try then
            sort schedules
        else
            try_one (counter+1) next_r_gen ((cost, this_sched):schedules)
            where
              this_sched = random_schedule dfg this_r_gen
              (dfg2,sched_map) = cost_estimate dfg this_sched
              cost = real_cost_estimate dfg2 sched_map
              (next_r_gen, this_r_gen) = R.split r_gen


-- |cost_estimate takes a DFG, a schedule for the operations in that DFG and
-- |initial and final queue states.  We simulate the execution of the schedule
-- |and estimate how much it would cost to make it queue legal.
-- |The input and output queues are embedded in the graph now.
cost_estimate :: Dfg -> [G.Node] -> (Dfg,IM.IntMap Int)
cost_estimate dfg_init schedule = (dfg_live_ins_actualized,complete_schedule)
  where
    -- Add scheduling edges to DFG
    -- Note: last_instr will be queue_back if there are no instrs
    (dfg_instrs_scheduled,last_instr) =
        foldl add_schedule_edge (dfg_init, queue_back) schedule
    add_schedule_edge (dfg_edging, prev) curr =
        (G.insEdge (prev, curr, DfgSeq) dfg_edging, curr)
    dfg_linked = G.insEdge (last_instr, out_queue_front, DfgSeq) dfg_instrs_scheduled

    -- Here we just follow the chain of next pointers in the DFG and assign numbers
    almost_complete_schedule = assign_schedule_numbers (Just queue_front) 0 IM.empty
    assign_schedule_numbers :: (Maybe G.Node) -> Int -> (IM.IntMap Int) -> (IM.IntMap Int)
    assign_schedule_numbers Nothing _ scheduling = scheduling
    assign_schedule_numbers (Just nid) sched scheduling =
        assign_schedule_numbers
            (dfg_next nid dfg_linked) (sched + 1) (IM.insert nid sched scheduling)
    complete_schedule =
        IM.insert live_out_hub_id (almost_complete_schedule IM.! out_queue_front)
              (IM.insert live_in_hub_id (almost_complete_schedule IM.! queue_back) almost_complete_schedule)

    -- A couple basic ordering functions for sorting live-ins/outs
    compare_op_and_port (op1_id,DfgPort p1) (op2_id,DfgPort p2) =
        if op1_id == op2_id then
            compare p1 p2
        else
            compare (complete_schedule IM.! op1_id) (complete_schedule IM.! op2_id)
    compare_op_and_port _ _ = error "not port!!!"

    compare_nets_by_schedule pre_not_suc (_,n1) (_,n2) =
        case (sortBy compare_op_and_port $ neighbors dfg_linked n1,
              sortBy compare_op_and_port $ neighbors dfg_linked n2) of
          (x1:_, x2:_) -> compare_op_and_port x1 x2
          (l1,l2) -> error ("none?!? "++show pre_not_suc++" "++show n1++":"++show l1++
                                          "  "++show n2++":"++show l2)
        where
          neighbors = if pre_not_suc then G.lpre else G.lsuc

    (actual_outs_unordered,_,_,formal_outs_unordered) = G.context dfg_linked live_out_hub_id
    (formal_ins_unordered ,_,_,actual_ins_unordered)  = G.context dfg_linked live_in_hub_id
    actual_outs_sorted = {--T.trace "compare 1"--} sortBy (compare_nets_by_schedule True ) actual_outs_unordered
    formal_outs_sorted = {--T.trace "compare 2"--} sortBy (compare_nets_by_schedule False) formal_outs_unordered
    actual_ins_sorted  = {--T.trace "compare 3"--} sortBy (compare_nets_by_schedule False) actual_ins_unordered
    formal_ins_sorted  = {--T.trace "compare 4"--} sortBy (compare_nets_by_schedule True ) formal_ins_unordered

    outs_zipped = assert (length actual_outs_sorted == length formal_outs_sorted)
                  zip actual_outs_sorted formal_outs_sorted
    dfg_live_outs_actualized = foldl live_out_link_up dfg_linked outs_zipped
    live_out_link_up dfg_live_outing ((_,actual_id), (_,formal_id)) =
      (ps,actual_id,label,filter (\(_,n) -> n /= live_out_hub_id) (ss1 ++ ss2))&dfg_no_formal
      where
        ((ps,_,label,ss1), dfg_no_actual) = g_unsafe_match ("rr34 ") actual_id dfg_live_outing
        ((_,_,_,ss2), dfg_no_formal)     = g_unsafe_match ("rr36 ") formal_id dfg_no_actual

    ins_zipped = assert (length actual_ins_sorted == length formal_ins_sorted)
                 zip formal_ins_sorted actual_ins_sorted
    ins_augmented = mapi augmentor ins_zipped
    augmentor i ((_,f_id),(_,a_id)) =
        (i, f_id, G.indeg  dfg_linked f_id,
         i, a_id, G.outdeg dfg_linked a_id)

    -- Okay... here is some tricky business.  We need to make sure that we pair
    -- up the formal live-ins and actual live-ins such that there are never more
    -- copies of a value in the queue than there are consumers for that value.
    -- [(a1,ap1,an1,f1,fp1,fn1),(a2,ap2,an2,f2,fp2,fn2),...]
    -- Sort according to: an_i > fn_i => bad_i; all bad before good
    --   among bad, sort by decreasing an_i
    -- Now we have the highest bad an_i first; need to find a swap
    --   swappable: fn_j >= an_i & an_j < an_i
    --   sort swappable: |ap_i-fp_j|+|ap_j-fp_i| (ascending)
    -- Swap! (a_i,ap_i,an_i,f_j,fp_j,fn_j):(a_j,ap_j,an_j,f_i,fp_i,fn_i):rest

    ins_rejigged = rejig_live_in_ordering ins_augmented
    rejig_live_in_ordering fixing =
      if is_bad_crap worst then
          rejig_live_in_ordering switched
      else
          -- 'sort' just works because the index we care about is the
          -- first member of the tuple
          sort fixing
      where
        is_bad_crap (_,_,n_formal,_,_,n_actual) = n_formal > n_actual
        worst_first_compare (p1@(_,_,n_formal1,_,_,_)) (p2@(_,_,n_formal2,_,_,_)) =
            case (is_bad_crap p1, is_bad_crap p2) of
              (False, False) -> EQ
              (True,  False) -> LT
              (False, True ) -> GT
              (True,  True ) -> compare n_formal2 n_formal1
        ((worst@(f_i_worst,f_worst,f_n_worst,a_i_worst,a_worst,a_n_worst)):better) =
            sortBy worst_first_compare fixing
        can_switch (_,_,f_n_better,_,_,a_n_better) =
            a_n_better >= f_n_worst && f_n_better < f_n_worst
        dist_from_worst (f_i,_,_,a_i,_,_) =
            abs (f_i_worst - a_i) + abs (a_i_worst - f_i)
        best_switch_compare p1 p2 =
            case (can_switch p1, can_switch p2) of
              (False, False) -> EQ
              (True,  False) -> LT
              (False, True ) -> GT
              (True,  True ) -> compare (dist_from_worst p1) (dist_from_worst p2)
        (f_i_switch,f_switch,f_n_switch,a_i_switch,a_switch,a_n_switch):rest =
            sortBy best_switch_compare better
        switched = 
            (f_i_worst,f_worst,f_n_worst,a_i_switch,a_switch,a_n_switch):
            (f_i_switch,f_switch,f_n_switch,a_i_worst,a_worst,a_n_worst):
            rest

    -- Man.  Live-ins suck so bad.  We're only half-way there.  Now we need to
    -- split up the nets appropriately.

    dfg_live_ins_actualized = case ins_zipped of
        [] -> dfg_live_outs_actualized
        _:_ -> foldl live_in_link_up dfg_live_outs_actualized ins_rejigged
    live_in_link_up dfg_live_inning ((_,orig_formal_net,_,_,orig_actual_net,_)) =
      distribute ps ss dfg_no_formal
      where
        ((_,_,label,ss), dfg_no_actual) =
            g_unsafe_match ("rr34 ") orig_actual_net dfg_live_inning
        ((ps,_,_,_), dfg_no_formal) =
            g_unsafe_match ("rr36 ") orig_formal_net dfg_no_actual

        distribute [] _ _ = error "dff"
        -- distribute [last_formal] last_actuals _ | T.trace ("foo "++show(length last_actuals)) False = undefined
        distribute [last_formal] last_actuals dfg_disting =
            ([last_formal],orig_actual_net,label,last_actuals)&dfg_disting
        distribute (formal:formals) actuals dfg_disting =
            assert (length actuals > length formals)
                   distribute formals actuals_remainder dfg_disted
            where
              this_fanout = (length actuals) `div_round` (1 + length formals)
              [net] = G.newNodes 1 dfg_disting
              (x,actuals_remainder) = splitAt this_fanout actuals
              dfg_disted = ([formal],net,label,x) & dfg_disting

-- okay, before we were just massaging the DFG into the proper form.  Now lets
-- really estimate the cost
real_cost_estimate :: Dfg -> IM.IntMap Int -> Int
real_cost_estimate dfg sched = pass_swap_cost1
  where
    -- A couple basic ordering functions for sorting live-ins/outs
    compare_op_and_port (DfgPort p1,op1_id) (DfgPort p2,op2_id) =
        if op1_id == op2_id then
            compare p1 p2
        else
            compare (sched IM.! op1_id) (sched IM.! op2_id)
    compare_op_and_port _ _ = error "not port!!!"

    -- XXX enhance filter when there are other kinds of nets
    queue_nets = filter (dfg_is_net . MB.fromJust . G.lab dfg) (G.nodes dfg)
    pass_swap_cost1 = foldl calculate_costs_on_net 0 queue_nets
    calculate_costs_on_net :: Int -> G.Node -> Int
    calculate_costs_on_net accumulated_cost outer_net =
      if consumer_count < 1 then
          T.trace ("drop!") accumulated_cost + 1 -- Count 1 for the drop
      else
          assert (latest_time > outer_driver_time)
          accumulated_cost + dup_cost + pass_swap_cost
      where
        dup_cost = -- T.trace ("dups: "++show (consumer_count - 1))
                   consumer_count - 1
        pass_swap_cost = --T.trace ("ps:"++show longest_chain_within_live_range_of_net++"  ss:"++show bar)
                         longest_chain_within_live_range_of_net + bar
        instrs_per_queue_cycle =
            (latest_time - outer_driver_time) `div_ceil` (longest_chain_within_live_range_of_net + 1)
        foo = (latest_time - outer_driver_time) `rem` instrs_per_queue_cycle
        bar = min foo (instrs_per_queue_cycle - foo)
        consumer_count = length consumers
        outer_driver_time = sched IM.! outer_driver
        latest_time = sched IM.! latest_consumer
        ([(DfgPort outer_driver_port,outer_driver)],_,_,consumers) = G.context dfg outer_net
        (DfgPort latest_port,latest_consumer) = maximumBy compare_op_and_port consumers
        longest_chain_within_live_range_of_net =
            longest_chain (IM.insert outer_driver 0 IM.empty) 0 (dfg_unsafe_next outer_driver dfg)
        -- longest_chain iterates over all the instructions within outer_net's
        -- live range in schedule-order.  Each invocation calculates the longest
        -- chain of instructions within the live range that end at that
        -- instruction.
        longest_chain :: IM.IntMap Int -> Int -> G.Node -> Int
        longest_chain longest_chain_ending_at longest_so_far inner_op =
          if inner_op == latest_consumer then
              -- We have scanned all the ops in outer_net's live range
              so_far'
          else
              -- traverse the next instruction in the schedule
              longest_chain ending_at' so_far' (dfg_unsafe_next inner_op dfg)
          where
            (ending_at',so_far') = foldl extend_chain
                          (IM.insert inner_op 0 longest_chain_ending_at,longest_so_far)
                          (filter data_ports_within_outer (G.lpre dfg inner_op))
            -- The need for data_ports_within_outer is a little subtle.  On the
            -- last instruction (latest_consumer), we want to make sure we
            -- account for exactly the ports before latest_port.
            data_ports_within_outer (_ {-XXX-} ,DfgPort pred_port) =
                sched IM.! inner_op < latest_time || pred_port < latest_port
            data_ports_within_outer (_,DfgSeq) = False
            extend_chain (m,l) (inner_net,_) =
                if inner_driver_time < outer_driver_time
                   || (inner_driver_time == outer_driver_time
                       && inner_driver_port < outer_driver_port) then
                    -- inner_net is not nested within outer_net's live-range
                    (m,l)
                else
                    (IM.insert inner_op new_longest_ending_at_inner m, new_longest)
                where
                  inner_driver_chain_length = m IM.! inner_driver
                  new_longest_ending_at_inner =
                      max (m IM.! inner_op) (inner_driver_chain_length + 1)
                  new_longest = max l new_longest_ending_at_inner
                  inner_driver_time = sched IM.! inner_driver
                  [(inner_driver,DfgPort inner_driver_port)] = G.lpre dfg inner_net

-- The rollouts are a little bit complicated because of the "real" instruction verus "queue"
-- instruction difference.

data UCT_Node = UCT_Node {
      uct_instr      :: Int,
      uct_sched      :: [Int],
      uct_the_q      :: Q.Seq Int,
      uct_ready_pool :: Q.Seq Int,
      uct_dfg        :: Dfg,
      uct_unexpanded :: Q.Seq Int,
      uct_incomplete :: Q.Seq UCT_Node,
      uct_complete   :: Q.Seq UCT_Node,
      uct_dead_end   :: Q.Seq UCT_Node,
      uct_vis_count  :: Integer,
      uct_estimate   :: Double }

roll_out :: R.RandomGen r => Dfg -> Int -> r -> UCT_Node
roll_out dfg instr rgen = UCT_Node 42 [] Q.empty Q.empty dfg Q.empty Q.empty Q.empty Q.empty 1 4.2

uct_descent :: R.RandomGen r => UCT_Node -> r -> UCT_Node
uct_descent parent rgen =
    case (Q.null parent_unexpanded, parent) of
      (False, _) -> roll_out (uct_dfg parent) instr rgen'
          where instr = Q.index parent_unexpanded ridx
                (ridx, rgen') = next rgen
      (True, _) -> roll_out (uct_dfg parent) 42 rgen
    where
      parent_unexpanded = uct_unexpanded parent
--       RollOut [] -> error "empty roll out?"
--       RollOut [just_one] -> parent
--       InternalNode instr [] -> parent
--       -- Interesting case #1: We have reached the frontier of the tree and need to do
--       -- roll-outs
--       RollOut _:instr_2:instr_n ->
--           parent {uct_child = InternalNode children 42}
--               where
--                 children = []
--     where
--       sched_p = uct_sched parent
--       the_q_p = uct_the_q parent
--       ready_pool_p = uct_ready_pool parent
--       dfg_p = uct_dfg blah
-- (Maybe (IntSeq, gr DfgNInfo DfgEInfo, IntSeq)) = try_instr instr1 the_q_prev dfg_prev ->
                
--           children = try_all dfg
--           try_all dfg = []
-- -- Note: Let's keep the children sorted by score
-- uct_descent (InternalNode instr (best:rest)) the_q ready_pool _ =
--     InternalNode instr children
--         where
--           children = (new_best:rest)
--           new_best = uct_descent best
--           -- XXX: sort children

-- unavfoo (UTreeNode (blah, [])) =
--   if done then (UTreeNode (blah, []))
--   else (UTreeNode (blah, []))
--   where done = True
-- unavfoo (UTreeNode (blah_orig, children_orig)) = UTreeNode (blah, children)
--   where
--     (child_orig, rest) = select children_orig
--     child = unavfoo child_orig
--     select x = (head x, tail x)
--     children = child:rest
--     blah = update blah_orig children
--     update x y = x



-- End DFG manipulation, scheduling, etc.



--
main :: IO ()
main = do
   progName <- Env.getProgName
   putStrLn ("The program name is: "++progName)
   args <- Env.getArgs
   let [ bbl_fname, sched_fname ] = args
   putStrLn bbl_fname
   putStrLn sched_fname
   bbl_string <- readFile bbl_fname
   let JS.Ok (JS.JSObject bbl_json) = JS.decode bbl_string
   let dfg = extract_bbl bbl_json
   --let (u1,u2) = sched_that_shit dfg

   SIO.withFile "file.dot" SIO.WriteMode (\h -> SIO.hPutStr h (Dot.showDot (Dot.fglToDot dfg)))
   -- Cmd.system("dot -Tpng -ofile.png file.dot")

   let schedules = random_schedule_try_a_few dfg (mkStdGen 42)
   let (best_cost,  best_sched):_  = schedules
   let (worst_cost, worst_sched):_ = reverse schedules
   --putStrLn ("Did random schedule "++show rrr)
   --let (dfg2,sched) = cost_estimate dfg rrr

   --SIO.withFile "file2.dot" SIO.WriteMode (\h -> SIO.hPutStr h (Dot.showDot (Dot.fglToDot dfg2)))
   -- Cmd.system("dot -Tpng -ofile.png file.dot")

   --let dsfdsf = real_cost_estimate dfg2 sched
   --let foo nid = (nid,g_unsafe_lab "gar" dfg nid)
   putStrLn ("Survived! ["++show best_cost++"  "++show best_sched++" -- "++
                          show worst_cost++"  "++show worst_sched++"]")
   writeFile sched_fname "XXX: TBD"
